#include "log.h"
#include "klib.h"
#include "mmu.h"
#include "mem.h"
#include "console.h"

extern uint32_t     kernel_base[], s_text[], e_text[], s_data[];    // from kernel.lds
extern uint32_t     mem_free_start;                                 // from kernel.lds
static addr_alloc_t paddr_alloc;
static uint32_t     total_memory_size;    // total mem size because of read mem

static pde_t kernel_page_dir[PDE_CNT] __attribute__((aligned(MEM_PAGE_SIZE)));

static void addr_alloc_init(addr_alloc_t *addr_alloc, uint8_t *bits, uint32_t start, uint32_t size, uint32_t page_size)
{
    mutex_init(&addr_alloc->mutex);
    addr_alloc->start     = start;
    addr_alloc->size      = size;
    addr_alloc->page_size = page_size;
    bitmap_init(&addr_alloc->bitmap, bits, addr_alloc->size / page_size, 0);
}

static uint32_t addr_alloc_page(addr_alloc_t *addr_alloc, int page_count)
{
    uint32_t addr = -1;
    mutex_lock(&addr_alloc->mutex);
    int page_idx = bitmap_alloc_nbits(&addr_alloc->bitmap, 0, page_count);
    if (page_idx >= 0)
    {
        addr = addr_alloc->start + page_idx * addr_alloc->page_size;
    }
    mutex_unlock(&addr_alloc->mutex);
    return addr;
}

static void addr_free_page(addr_alloc_t *addr_alloc, uint32_t addr, int page_count)
{
    mutex_lock(&addr_alloc->mutex);
    uint32_t page_idx = (addr - addr_alloc->start) / addr_alloc->page_size;
    bitmap_set_bit(&addr_alloc->bitmap, page_idx, page_count, 0);
    mutex_unlock(&addr_alloc->mutex);
}

static void show_mem_info(boot_info_t *boot_info)
{
    log_printf("mem region:");
    for (int i = 0; i < boot_info->ram_region_count; i++)
    {
        log_printf("[%d],start:%x, size:%x", i, boot_info->ram_region_cfg[i].start, boot_info->ram_region_cfg[i].size);
    }
}
static uint32_t total_mem_size(boot_info_t *boot_info)
{
    uint32_t mem_size = 0;
    for (int i = 0; i < boot_info->ram_region_count; i++)
    {
        mem_size += boot_info->ram_region_cfg[i].size;
    }
    return mem_size;
}
pte_t *find_pte(pde_t *page_dir, uint32_t vaddr, int alloc)
{
    pte_t *page_table;
    pde_t *pde = page_dir + pde_index(vaddr);
    if (pde->present)
    {
        page_table = (pte_t *)pde_paddr(pde);
    }
    else
    {
        if (!alloc)
        {
            return (pte_t *)0;
        }
        uint32_t pg_addr = addr_alloc_page(&paddr_alloc, 1);
        if (pg_addr == -1)
        {
            return (pte_t *)0;
        }
        pde->v =
            pg_addr | PDE_P | PDE_W |
            PDE_U;    // this pg_addr is the real 32bit address, but only 20bit is used in pde, so can do in this way
        page_table = (pte_t *)pg_addr;
        kernel_memset(page_table, 0, MEM_PAGE_SIZE);
    }
    uint32_t pte_idx = pte_index(vaddr);
    page_table += pte_idx;    // this is the add of pointer, so plus 4 rather than 1
    return page_table;
}
// get paddr from vaddr
uint32_t mem_get_paddr(uint32_t page_dir, uint32_t vaddr)
{
    pte_t *pte = find_pte((pde_t *)page_dir, vaddr, 0);
    if (pte == (pte_t *)0)
    {
        return 0;
    }
    return pte_paddr(pte) + (vaddr & (MEM_PAGE_SIZE - 1));
}
static int mem_create_map(pde_t *page_dir, uint32_t vaddr, uint32_t paddr, uint32_t count, uint32_t perm)
{
    for (int i = 0; i < count; i++)
    {
        // log_printf("create map: vaddr: %x, paddr: %x perm: %x",vaddr,paddr,perm);
        pte_t *pte = find_pte(page_dir, vaddr, 1);
        if (pte == (pte_t *)0)
        {
            log_printf("find_pte failed!");
            return -1;
        }
        // log_printf("pte addr: %x",(uint32_t)pte);
        ASSERT(pte->present == 0);
        pte->v = paddr | perm | PTE_P;
        vaddr += MEM_PAGE_SIZE;
        paddr += MEM_PAGE_SIZE;
    }
}
void create_kernel_table(void)
{
    static mem_map_t kernel_map[] = {
        {kernel_base, s_text, kernel_base, PTE_W},          // 0-0x10000, pg_count=16
        {s_text, e_text, s_text, 0},                        // 0x10000-0x14000,pa_count=4
        {s_data, (void *)MEM_EBDA_START, s_data, PTE_W},    // 0x14000-0x80000,pg_count=108
        {(void *)CONSOLE_DISP_ADDR, (void *)CONSOLE_DISP_END, (void *)CONSOLE_DISP_ADDR, PTE_W},
        // {(void*)MEM_EXT_START,(void*)(down2(total_memory_size,MEM_PAGE_SIZE)),(void*)MEM_EXT_START,PTE_W | PTE_U},
        {(void *)MEM_EXT_START, (void *)MEM_EXT_END, (void *)MEM_EXT_START, PTE_W},
    };
    // I think there a problems with the border between 2 mem_map_t
    for (int i = 0; i < sizeof(kernel_map) / sizeof(mem_map_t); i++)
    {
        mem_map_t *map        = kernel_map + i;
        uint32_t   vstart     = down2((uint32_t)map->vstart, MEM_PAGE_SIZE);
        uint32_t   vend       = up2((uint32_t)map->vend, MEM_PAGE_SIZE);
        uint32_t   pstart     = down2((uint32_t)map->pstart, MEM_PAGE_SIZE);
        int        page_count = (vend - vstart) / MEM_PAGE_SIZE;

        mem_create_map(kernel_page_dir, vstart, pstart, page_count, map->perm);
    }
}

void mem_init(boot_info_t *boot_info)
{
    uint8_t *mem_free = (uint8_t *)&mem_free_start;
    log_printf("mem init:");
    show_mem_info(boot_info);
    uint32_t total_size = total_mem_size(boot_info);
    total_memory_size   = total_size;
    log_printf("mem size: %d", total_size);
    uint32_t mem_up1MB_free = total_size - MEM_EXT_START;
    mem_up1MB_free          = down2(mem_up1MB_free, MEM_PAGE_SIZE);
    log_printf("free memory (aligned to page size): start: %x, size: %x", MEM_EXT_START, mem_up1MB_free);

    addr_alloc_init(&paddr_alloc, mem_free, MEM_EXT_START, mem_up1MB_free, MEM_PAGE_SIZE);
    mem_free += bitmap_byte_count(paddr_alloc.size / MEM_PAGE_SIZE);
    ASSERT(mem_free < (uint8_t *)MEM_EBDA_START);

    create_kernel_table();
    mmu_set_page_dir((uint32_t)kernel_page_dir);
}
uint32_t mem_create_uvm(void)
{
    uint32_t paddr = addr_alloc_page(&paddr_alloc, 1);
    if (paddr == -1)
    {
        return 0;
    }
    pde_t *page_dir = (pde_t *)paddr;
    // create identity mapping when vaddr < MEM_PAGE_SIZE
    kernel_memset((void *)page_dir, 0, MEM_PAGE_SIZE);
    // create pte for (0x0 ~ 0x80000000)
    uint32_t user_pde_idx_start = pde_index(MEM_TASK_BASE);
    for (int i = 0; i < user_pde_idx_start; i++)
    {
        page_dir[i].v = kernel_page_dir[i].v;
    }
    return (uint32_t)page_dir;
}

void mem_init_test(boot_info_t *boot_info)
{
    addr_alloc_t addr_alloc;
    uint8_t      bits[8];
    addr_alloc_init(&addr_alloc, bits, 0x1000, 64 * 4096, 4096);
    for (int i = 0; i < 32; i++)
    {
        uint32_t addr = addr_alloc_page(&addr_alloc, 2);
        log_printf("alloc addr:%x", addr);
    }

    uint32_t addr = 0x1000;
    for (int i = 0; i < 32; i++)
    {
        addr_free_page(&addr_alloc, addr, 2);
        addr += 2 * 4096;
    }
}
int mem_alloc_for_page_dir(uint32_t page_dir, uint32_t vaddr, uint32_t size, int perm)
{
    uint32_t curr_vaddr = vaddr;
    int      page_count = up2(size, MEM_PAGE_SIZE) / MEM_PAGE_SIZE;
    for (int i = 0; i < page_count; i++)
    {
        uint32_t paddr = addr_alloc_page(&paddr_alloc, 1);
        if (paddr == -1)
        {
            log_printf("mem_alloc_for_page_dir: alloc_failed. no memory");
            return -1;
        }
        int err = mem_create_map((pde_t *)page_dir, curr_vaddr, paddr, 1, perm);
        if (err < 0)
        {
            log_printf("mem_alloc_for_page_dir: mem_create_map failed err:%d", err);
            // in fact, this condition should free the page unused
            return -1;
        }
        curr_vaddr += MEM_PAGE_SIZE;
    }
    return 0;
}
int mem_alloc_page_for(uint32_t vaddr, uint32_t size, int perm)
{
    return mem_alloc_for_page_dir(task_current_task()->tss.cr3, vaddr, size, perm);
}
uint32_t mem_alloc_page(void)
{
    uint32_t paddr = addr_alloc_page(&paddr_alloc, 1);
    if (paddr == -1)
    {
        return 0;
    }
    return paddr;
}
static pde_t *curr_page_dir(void)
{
    return (pde_t *)task_current_task()->tss.cr3;
}
void mem_free_page(uint32_t addr)
{
    if (addr < MEM_TASK_BASE)
    {
        addr_free_page(&paddr_alloc, addr, 1);
    }
    else
    {
        pte_t *pte = find_pte(curr_page_dir(), addr, 0);
        ASSERT(pte != (pte_t *)0 && pte->present);
        addr_free_page(&paddr_alloc, pte_paddr(pte), 1);
        pte->v = 0;
    }
}
void mem_destroy_uvm(uint32_t page_dir)
{
    ASSERT(page_dir != 0);
    uint32_t user_pde_start = pde_index(MEM_TASK_BASE);
    pde_t   *pde            = (pde_t *)page_dir + user_pde_start;
    for (int i = user_pde_start; i < PDE_CNT; i++, pde++)
    {
        if (!pde->present)
        {
            continue;
        }
        pte_t *pte = (pte_t *)pde_paddr(pde);
        for (int j = 0; j < PTE_CNT; j++, pte++)
        {
            if (!pte->present)
            {
                continue;
            }
            addr_free_page(&paddr_alloc, pte_paddr(pte), 1);
        }
        addr_free_page(&paddr_alloc, (uint32_t)pde_paddr(pde), 1);
    }
    addr_free_page(&paddr_alloc, page_dir, 1);
}
uint32_t mem_copy_uvm(uint32_t page_dir)
{
    uint32_t to_page_dir = mem_create_uvm();
    if (to_page_dir == 0)
    {
        goto mem_copy_uvm_failed;
    }
    uint32_t user_pde_start = pde_index(MEM_TASK_BASE);
    pde_t   *pde            = (pde_t *)page_dir + user_pde_start;
    for (int i = user_pde_start; i < PDE_CNT; i++, pde++)
    {
        if (!pde->present)
        {
            continue;
        }
        pte_t *pte = (pte_t *)pde_paddr(pde);
        for (int j = 0; j < PTE_CNT; j++, pte++)
        {
            if (!pte->present)
            {
                continue;
            }
            uint32_t page_paddr = addr_alloc_page(&paddr_alloc, 1);
            if (page_paddr == -1)
            {
                goto mem_copy_uvm_failed;
            }
            uint32_t vaddr = (i << 22) | (j << 12);
            int      err   = mem_create_map((pde_t *)to_page_dir, vaddr, page_paddr, 1, get_pte_perm(pte));
            if (err < 0)
            {
                goto mem_copy_uvm_failed;
            }
            kernel_memcpy((void *)page_paddr, (void *)vaddr, MEM_PAGE_SIZE);
        }
    }
    return to_page_dir;
mem_copy_uvm_failed:
    if (to_page_dir)
    {
        mem_destroy_uvm(to_page_dir);
    }
    return -1;
}
int mem_copy_uvm_data(uint32_t to, uint32_t page_dir, uint32_t from, int size)
{
    while (size > 0)
    {
        uint32_t to_paddr = mem_get_paddr(page_dir, to);
        if (!to_paddr)
        {
            return -1;
        }
        // aim to calc the size need to copy
        uint32_t offset_in_page = to_paddr & (MEM_PAGE_SIZE - 1);
        uint32_t curr_size      = MEM_PAGE_SIZE - offset_in_page;
        if (curr_size >= size)
        {
            curr_size = size;
        }
        kernel_memcpy((void *)to_paddr, (void *)from, curr_size);
        size -= curr_size;
        to += curr_size;
        from += curr_size;
    }
    return 0;
}
char *sys_sbrk(int incr)
{
    ASSERT(incr >= 0);
    int     pre_incr     = incr;
    task_t *task         = task_current_task();
    char   *pre_heap_end = (char *)task->heap_end;
    if (incr == 0)
    {
        log_printf("sbrk(0):end=%x", pre_heap_end);
        return pre_heap_end;
    }
    uint32_t start = task->heap_end;
    uint32_t end   = task->heap_end + incr;

    // 0x80001024 - 0x80001048
    int start_offset = start % MEM_PAGE_SIZE;    // 0x24
    if (start_offset)
    {
        if (start_offset + incr <= MEM_PAGE_SIZE)
        {
            task->heap_end = end;
            log_printf("sbrk(%d):end=%x", pre_incr, end);
            return pre_heap_end;
        }
        else
        {
            uint32_t curr_size = MEM_PAGE_SIZE - start_offset;
            start += curr_size;
            incr -= curr_size;
        }
    }
    if (incr)
    {
        int err = mem_alloc_page_for(start, incr, PTE_P | PTE_W | PTE_U);
        if (err < 0)
        {
            log_printf("sbrk mem_alloc_page_for failed");
            return (char *)-1;
        }
    }
    log_printf("sbrk(%d):end=%x", pre_incr, end);
    task->heap_end = end;
    return (char *)pre_heap_end;
}